package com.yanghui.dtracing.plugin.springmvc;

import com.yanghui.dtracing.core.constant.AgentConstant;
import com.yanghui.dtracing.core.constant.TypeEnum;
import com.yanghui.dtracing.core.opentracing.*;
import com.yanghui.dtracing.core.opentracing.log.LogData;
import com.yanghui.dtracing.core.plugin.interceptor.AbstractInstanceMethodInterceptor;
import com.yanghui.dtracing.core.plugin.interceptor.MethodInterceptResult;
import io.opentracing.Span;
import io.opentracing.propagation.Format;
import io.opentracing.propagation.TextMapExtract;
import io.opentracing.tag.Tags;
import lombok.AllArgsConstructor;

import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletResponse;
import java.io.ByteArrayInputStream;
import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
import java.lang.reflect.Method;
import java.util.HashMap;
import java.util.Iterator;
import java.util.Map;

public class SpringMvcMethodAroundInterceptor extends AbstractInstanceMethodInterceptor {

    @AllArgsConstructor
    private static class HttpHeaderCarrier implements TextMapExtract {
        private HttpServletRequest request;
        @Override
        public Iterator<Map.Entry<String, String>> iterator() {
            Map<String,String> map = new HashMap<>();
            map.put(AgentConstant.HTTP_HEADER_SPAN_CONTEXT_KEY,request.getHeader(AgentConstant.HTTP_HEADER_SPAN_CONTEXT_KEY));
            return map.entrySet().iterator();
        }
    }

    @Override
    public boolean isIntercept(Object thisObj, Method method, Object[] allArguments) throws Throwable {
        HttpServletRequest req = (HttpServletRequest)allArguments[0];
        if(req.getRequestURI().contains("hystrix.stream")
                || req.getRequestURI().contains("prometheus")
                || req.getRequestURI().contains("metrics")
                || req.getRequestURI().contains("favicon")
                || req.getRequestURI().contains("healthCheck")
                || req.getRequestURI().contains("app/dependency/endpoint")
                || req.getRequestURI().contains("simple/health/check")
                || req.getRequestURI().contains("bycdao-ui")
                || req.getRequestURI().contains("doc.html")
                || req.getRequestURI().contains("/topology")
                || req.getRequestURI().contains("swagger")
                || req.getRequestURI().contains("/job/endpoint")
                || req.getRequestURI().contains("/gracefulShutdown")
                || req.getRequestURI().contains("/jvmInfo")
                || req.getRequestURI().contains("/loggers")
                || req.getRequestURI().contains("/topology")){
            return false;
        }
        return true;
    }

    @Override
    protected void doBuild(DTracer.SpanBuilder spanBuilder, Object thisObj, Method method, Object[] allArguments, MethodInterceptResult result) {
        HttpServletRequest request = (HttpServletRequest)allArguments[0];
        DTracer tracer = Tracing.getTracer();
        /**恢复上下文**/
        DTracerSpanContext context = (DTracerSpanContext)tracer.extract(Format.Builtin.TEXT_MAP_EXTRACT,new HttpHeaderCarrier(request));
        spanBuilder.withTag(Tags.COMPONENT.getKey(), DTracerTags.COMPONENT_HTTP)
                .withTag(DTracerTags.TYPE.getKey(), TypeEnum.ENTER.getName())
                .withTag(Tags.HTTP_URL.getKey(),request.getRequestURL().toString())
                .withTag(DTracerTags.HTTP_PATH,request.getRequestURI())
                .withTag(Tags.HTTP_METHOD.getKey(),request.getMethod()).ignoreActiveSpan();
        if(null != context){
            String autoId = context.getBaggageItem(Baggage.AUTO_ID_KEY);
            if(autoId != null && !"".equals(autoId)){
                context.addBaggageItem(Baggage.AUTO_ID_KEY,String.valueOf(Integer.valueOf(autoId)+ 100000));
            }
            spanBuilder.asChildOf(context);
        }
    }

    @Override
    protected void doBeforeMethod(Object thisObj, Method method, Object[] allArguments, MethodInterceptResult result, Span span) throws Throwable {
        span.log(System.currentTimeMillis(), LogData.SERVER_RECV_EVENT_VALUE);
        HttpServletResponse response = (HttpServletResponse)allArguments[1];
        HttpResponseWrapper httpResponseWrapper = new HttpResponseWrapper(response);
        allArguments[1] = httpResponseWrapper;
        result.getDeliveryParams().put(HttpServletResponse.class.getSimpleName(),response);
    }

    @Override
    protected void doAfterMethod(Object thisObj, Method method, Object[] allArguments, MethodInterceptResult result, Span span) throws Throwable {
        HttpResponseWrapper wrapper = (HttpResponseWrapper)allArguments[1];
        span.setTag(Tags.HTTP_STATUS.getKey(),Integer.valueOf(wrapper.getStatus()))
                .log(System.currentTimeMillis(),LogData.SERVER_SEND_EVENT_VALUE);
        if(wrapper.getStatus() == 500){
            span.setTag(Tags.ERROR.getKey(),true);
        }
        byte[] data = wrapper.getContent();
        String resultStr = new String(data,"utf-8");
        result.setResult(resultStr);
        HttpServletResponse response = (HttpServletResponse)result.getDeliveryParams().get(HttpServletResponse.class.getSimpleName());
        OutputStream outputStream = response.getOutputStream();
        try{
            writeResponse(new ByteArrayInputStream(data),outputStream);
        }finally {
            outputStream.flush();
        }
    }

    private void writeResponse(InputStream zin, OutputStream out) throws IOException {
        byte[] bytes = new byte[8192];
        int bytesRead = -1;
        while ((bytesRead = zin.read(bytes)) != -1) {
            out.write(bytes, 0, bytesRead);
        }
    }
}
